from pathlib import Path

from ipyniivue import download_dataset

BASE_API_URL = "https://niivue.com/demos/images"
DATA_FOLDER = Path("images")

# Download data for example
download_dataset(
    BASE_API_URL,
    DATA_FOLDER,
    files=[
        "mni152.nii.gz",
        "BrainMesh_ICBM152.lh.mz3",
    ],
)

mni152.nii.gz already exists.
BrainMesh_ICBM152.lh.mz3 already exists.
Dataset downloaded successfully to images.

import ipywidgets as widgets

from ipyniivue import DragMode, NiiVue, ShowRender, SliceType

# UI Widgets

dpi_check = widgets.Checkbox(value=False, description="HighDPI")
zoom_check = widgets.Checkbox(value=False, description="2D/3D Zoom Sync")

slice_type = widgets.Dropdown(
    options=[
        ("Axial", SliceType.AXIAL),
        ("Coronal", SliceType.CORONAL),
        ("Sagittal", SliceType.SAGITTAL),
        ("Render", SliceType.RENDER),
        ("A+C+S+R", SliceType.MULTIPLANAR),
    ],
    value=SliceType.MULTIPLANAR,
    description="View",
)

controller = widgets.Dropdown(
    options=[
        ("Independent", 0),
        ("Top Controls Bottom", 1),
        ("Bottom Controls Top", 2),
        ("Bidirectional", 3),
    ],
    value=1,
    description="Broadcast",
)

drag_mode = widgets.Dropdown(
    options=[
        ("contrast", DragMode.CONTRAST),
        ("measurement", DragMode.MEASUREMENT),
        ("pan/zoom", DragMode.PAN),
        ("none", DragMode.NONE),
    ],
    value=DragMode.CONTRAST,
    description="Drag",
)

mesh_clip_slider = widgets.FloatSlider(
    value=11, min=0, max=11.1, step=0.1, description="Mesh Clip"
)

# Info labels
intensity_label_1 = widgets.Label(value=" ")
intensity_label_2 = widgets.Label(value=" ")

# Container for controls
ui_controls = widgets.VBox(
    [
        widgets.HBox([dpi_check, zoom_check]),
        widgets.HBox([slice_type, controller]),
        widgets.HBox([drag_mode, mesh_clip_slider]),
    ]
)

# Initialize NiiVue Instances

nv1 = NiiVue()
nv2 = NiiVue()

# Configure NV1
nv1.opts.back_color = (1, 1, 1, 1)
nv1.opts.multiplanar_show_render = ShowRender.ALWAYS
nv1.opts.is_slice_mm = True
nv1.opts.is_orient_cube = True
nv1.set_clip_plane(0, 180, 40)

# Configure NV2
nv2.opts.back_color = (1, 1, 1, 1)
nv2.opts.multiplanar_show_render = ShowRender.ALWAYS
nv2.opts.is_slice_mm = True
nv2.set_clip_plane(0, 180, 40)
nv2.set_high_resolution_capable(True)

# Load Data
vol_data = [{"path": DATA_FOLDER / "mni152.nii.gz"}]
mesh_data_1 = [
    {"path": DATA_FOLDER / "BrainMesh_ICBM152.lh.mz3", "rgba255": [212, 212, 255, 255]}
]
mesh_data_2 = [
    {"path": DATA_FOLDER / "BrainMesh_ICBM152.lh.mz3", "rgba255": [222, 255, 222, 255]}
]

nv1.load_volumes(vol_data)
nv1.load_meshes(mesh_data_1)
nv2.load_volumes(vol_data)
nv2.load_meshes(mesh_data_2)


def on_mesh_loaded(mesh):
    """Set mesh shader."""
    nv1.set_mesh_shader(mesh.id, "Outline")


nv1.on_mesh_loaded(on_mesh_loaded)

# Sync Logic Init
sync_settings = {"3d": True, "2d": True, "clip_plane": True}
nv1.broadcast_to(nv2, sync_settings)  # Default: Top Controls Bottom

# Event Handlers


def on_slice_type_change(change):
    """Handle sliceType change."""
    nv1.set_slice_type(change["new"])
    nv2.set_slice_type(change["new"])


def on_controller_change(change):
    """Handle controller change."""
    val = change["new"]
    # 0: Independent
    if val == 0:
        nv1.broadcast_to([])
        nv2.broadcast_to([])
    # 1: Top Controls Bottom
    elif val == 1:
        nv1.broadcast_to(nv2, sync_settings)
        nv2.broadcast_to([])
    # 2: Bottom Controls Top
    elif val == 2:
        nv1.broadcast_to([])
        nv2.broadcast_to(nv1, sync_settings)
    # 3: Bidirectional
    elif val == 3:
        nv1.broadcast_to(nv2, sync_settings)
        nv2.broadcast_to(nv1, sync_settings)


def on_drag_mode_change(change):
    """Handle dragMode change."""
    nv1.opts.drag_mode = change["new"]
    nv2.opts.drag_mode = change["new"]


def on_dpi_change(change):
    """Handle dpiCheck change."""
    nv1.set_high_resolution_capable(change["new"])


def on_zoom_sync_change(change):
    """Handle zoomCheck change."""
    is_checked = change["new"]
    nv1.opts.yoke_3d_to_2d_zoom = is_checked
    nv2.opts.yoke_3d_to_2d_zoom = is_checked
    if is_checked:
        if len(nv1.scene.pan2d_xyzmm) > 3:
            nv1.scene.vol_scale_multiplier = nv1.scene.pan2d_xyzmm[3]
            nv1.draw_scene()
        if len(nv2.scene.pan2d_xyzmm) > 3:
            nv2.scene.vol_scale_multiplier = nv2.scene.pan2d_xyzmm[3]
            nv2.draw_scene()


def on_mesh_clip_change(change):
    """Handle mesh clip change."""
    val = change["new"]
    if val > 10:
        val = float("inf")
    nv1.opts.mesh_thickness_on_2d = val


def handle_location_change_1(data):
    """Handle nv1 location change."""
    intensity_label_1.value = data["string"]


def handle_location_change_2(data):
    """Handle nv2 location change."""
    intensity_label_2.value = data["string"]


# Attach observers
slice_type.observe(on_slice_type_change, names="value")
controller.observe(on_controller_change, names="value")
drag_mode.observe(on_drag_mode_change, names="value")
dpi_check.observe(on_dpi_change, names="value")
zoom_check.observe(on_zoom_sync_change, names="value")
mesh_clip_slider.observe(on_mesh_clip_change, names="value")

nv1.on_location_change(handle_location_change_1)
nv2.on_location_change(handle_location_change_2)

# Display

widgets.VBox(
    [
        ui_controls,
        widgets.VBox([nv1, nv2]),
        widgets.HBox([intensity_label_1, intensity_label_2]),
    ]
)